//===----------------------------------------------------------------------===//
//
//                         BusTub
//
// aggregation_executor.cpp
//
// Identification: src/execution/aggregation_executor.cpp
//
// Copyright (c) 2015-2021, Carnegie Mellon University Database Group
//
//===----------------------------------------------------------------------===//
#include <memory>
#include <vector>

#include "execution/executors/aggregation_executor.h"

namespace bustub {

AggregationExecutor::AggregationExecutor(ExecutorContext *exec_ctx, const AggregationPlanNode *plan,
                                         std::unique_ptr<AbstractExecutor> &&child)
    : AbstractExecutor(exec_ctx), plan_(plan), child_(std::move(child)) {}

void AggregationExecutor::Init() {
  child_->Init();
  aht_ = std::make_unique<SimpleAggregationHashTable>(plan_->GetAggregates(), plan_->GetAggregateTypes());
  aggregated_ = false;
  is_empty_ = false;
  Tuple ch_tuple;
  RID ch_rid;
  while (child_->Next(&ch_tuple, &ch_rid)) {
    aht_->InsertCombine(MakeAggregateKey(&ch_tuple), MakeAggregateValue(&ch_tuple));
  }
  if (aht_->Begin() == aht_->End()) {  // set initvalue
    is_empty_ = true;
  }
  aht_iterator_ = std::make_unique<SimpleAggregationHashTable::Iterator>(aht_->Begin());
}

auto AggregationExecutor::Next(Tuple *tuple, RID *rid) -> bool {
  if (aggregated_) {
    return false;
  }
  if (is_empty_) {
    std::vector<Value> vals;
    auto &group_bys = plan_->GetGroupBys();
    if (group_bys.empty()) {
      const auto &aggregates = aht_->GenerateInitialAggregateValue().aggregates_;
      vals.reserve(aggregates.size());
      vals.insert(vals.end(), aggregates.begin(), aggregates.end());
      Tuple result(vals, &GetOutputSchema());
      *tuple = std::move(result);
    } else {
      aggregated_ = true;
      return false;
    }

    aggregated_ = true;
    return true;
  }
  while (*aht_iterator_ != aht_->End()) {
    std::vector<Value> vals;
    auto key = aht_iterator_->Key();
    auto val = aht_iterator_->Val();
    vals.reserve(key.group_bys_.size() + val.aggregates_.size());
    vals.insert(vals.end(), key.group_bys_.begin(), key.group_bys_.end());
    vals.insert(vals.end(), val.aggregates_.begin(), val.aggregates_.end());
    Tuple result(vals, &GetOutputSchema());
    *tuple = std::move(result);
    ++*aht_iterator_;
    return true;
  }
  aggregated_ = true;
  return false;
}

auto AggregationExecutor::GetChildExecutor() const -> const AbstractExecutor * { return child_.get(); }

}  // namespace bustub
